package robotics.sandbox;

import java.awt.Transparency;
import java.awt.color.ColorSpace;
import java.awt.image.BufferedImage;
import java.awt.image.ComponentColorModel;
import java.awt.image.DataBuffer;
import java.awt.image.DataBufferByte;
import java.awt.image.PixelInterleavedSampleModel;
import java.awt.image.Raster;
import java.awt.image.WritableRaster;

/**
 * @author Asee Shrestha
 *
 */
/**
 * Thresholds the image at the given ThresholdValue
 * @params requiredThresholdValue The threshold value by which the image is to be threshold
 */
public class ImageOperations {

    //<editor-fold defaultstate="collapsed" desc="OtsuThresholder">
    private static class OtsuThresholder {

        private int histData[];
        private int maxLevelValue;
        private int threshold;

        public OtsuThresholder() {
            histData = new int[256];
        }

        public int[] getHistData() {
            return histData;
        }

        public int getMaxLevelValue() {
            return maxLevelValue;
        }

        public int getThreshold() {
            return threshold;
        }

        public int doThreshold(byte[] srcData, byte[] monoData) {
            int ptr;

            // Clear histogram data
            // Set all values to zero
            ptr = 0;
            while (ptr < histData.length) {
                histData[ptr++] = 0;
            }

            // Calculate histogram and find the level with the max value
            // Note: the max level value isn't required by the Otsu method
            ptr = 0;
            maxLevelValue = 0;
            while (ptr < srcData.length) {
                int h = 0xFF & srcData[ptr];
                histData[h]++;
                if (histData[h] > maxLevelValue) {
                    maxLevelValue = histData[h];
                }
                ptr++;
            }

            // Total number of pixels
            int total = srcData.length;

            float sum = 0;
            for (int t = 0; t < 256; t++) {
                sum += t * histData[t];
            }

            float sumB = 0;
            int wB = 0;
            int wF = 0;

            float varMax = 0;
            threshold = 0;

            for (int t = 0; t < 256; t++) {
                wB += histData[t];					// Weight Background
                if (wB == 0) {
                    continue;
                }

                wF = total - wB;						// Weight Foreground
                if (wF == 0) {
                    break;
                }

                sumB += (float) (t * histData[t]);

                float mB = sumB / wB;				// Mean Background
                float mF = (sum - sumB) / wF;		// Mean Foreground

                // Calculate Between Class Variance
                float varBetween = (float) wB * (float) wF * (mB - mF) * (mB - mF);

                // Check if new maximum found
                if (varBetween > varMax) {
                    varMax = varBetween;
                    threshold = t;
                }
            }

            // Apply threshold to create binary image
            if (monoData != null) {
                ptr = 0;
                while (ptr < srcData.length) {
                    monoData[ptr] = ((0xFF & srcData[ptr]) >= threshold) ? (byte) 255 : 0;
                    ptr++;
                }
            }

            return threshold;
        }
    }
    //</editor-fold>

    //<editor-fold defaultstate="collapsed" desc="public API">
    public static BufferedImage threshold(BufferedImage img, int requiredThresholdValue) {

        int height = img.getHeight();
        int width = img.getWidth();
        BufferedImage finalThresholdImage = new BufferedImage(width, height, BufferedImage.TYPE_INT_RGB);

        int red = 0;
        int green = 0;
        int blue = 0;

        for (int x = 0; x < width; x++) {
//			System.out.println("Row: " + x);
            try {

                for (int y = 0; y < height; y++) {
                    int color = img.getRGB(x, y);

                    red = ImageOperations.getRed(color);
                    green = ImageOperations.getGreen(color);
                    blue = ImageOperations.getBlue(color);

//					System.out.println("threshold : " + requiredThresholdValue);
                    if ((red + green + blue) / 3 < (int) (requiredThresholdValue)) {
                        finalThresholdImage.setRGB(x, y, ImageOperations.mixColor(0, 0, 0));
                    } else {
                        finalThresholdImage.setRGB(x, y, ImageOperations.mixColor(255, 255, 255));
                    }

                }
            } catch (Exception e) {
                e.getMessage();
            }
        }

        return finalThresholdImage;
    }

    public static BufferedImage thresholdOtsu(BufferedImage img) {
        OtsuThresholder thr = new OtsuThresholder();
        DataBufferByte db = (DataBufferByte) img.getData().getDataBuffer();

        byte[] srcData = db.getData(0);
        byte[] dstData = new byte[srcData.length];
        int threshold = thr.doThreshold(srcData, dstData);
        return threshold(img, threshold);
    }

    public static void drawFrame(BufferedImage img) {
        GreyFrame.drawFrame(img);
    }
    //</editor-fold>

    //<editor-fold defaultstate="collapsed" desc="utils">
    public static int mixColor(int red, int green, int blue) {
        return red << 16 | green << 8 | blue;
    }

    public static int getRed(int color) {
        return (color & 0x00ff0000) >> 16;
    }

    public static int getGreen(int color) {
        return (color & 0x0000ff00) >> 8;
    }

    public static int getBlue(int color) {
        return (color & 0x000000ff) >> 0;
    }

    public static BufferedImage mkImage(int width, int height, byte[] data) {
        DataBufferByte dataBuffer = new DataBufferByte(data, data.length, 0);

        PixelInterleavedSampleModel sampleModel =
                new PixelInterleavedSampleModel(DataBuffer.TYPE_BYTE, width, height, 1, width, new int[]{0});
        ColorSpace colourSpace = ColorSpace.getInstance(ColorSpace.CS_GRAY);
        ComponentColorModel colourModel = new ComponentColorModel(
                colourSpace, new int[]{8}, false, false, Transparency.OPAQUE, DataBuffer.TYPE_BYTE);

        WritableRaster raster = Raster.createWritableRaster(sampleModel, dataBuffer, null);

        return new BufferedImage(colourModel, raster, false, null);
    }
    //</editor-fold>
}
